"""风险控制 Agent."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Optional

from ..prompts import build_risk_control_prompt, load_task_template
from .base import AgentConfig, BaseETFAnalysisAgent


@dataclass
class RiskControlContext:
    fund: str
    current_time: str
    history_summary: str
    web_sources: str
    quant_metrics: str


class RiskControlAgent(BaseETFAnalysisAgent):
    """调用 LLM 生成风险控制报告."""

    def __init__(self, *, config: Optional[AgentConfig] = None) -> None:
        super().__init__(config=config)
        self._prompt = build_risk_control_prompt()
        self._task_template = load_task_template("risk_control_officer_task")

    def build_inputs(self, context: RiskControlContext) -> Dict[str, str]:
        return {
            "current_time": context.current_time,
            "fund": context.fund,
            "history_summary": context.history_summary,
            "web_sources": context.web_sources,
            "quant_metrics": context.quant_metrics,
            "task_description": self._task_template.description,
        }

    def run(self, context: RiskControlContext) -> str:
        inputs = self.build_inputs(context)
        formatted = self._prompt.format(**inputs)
        return self.llm.invoke(formatted).content  # type: ignore[return-value]

